package com.quanxiaoha.ai.robot.controller;

import jakarta.annotation.Resource;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.deepseek.DeepSeekAssistantMessage;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author Caius
 * @description DeepSeek R1 推理模型聊天
 * @since Created in 2025-09-04
 */
@RestController
@RequestMapping("/v1/ai")
public class DeepSeekR1ChatController {
    @Resource
    private DeepSeekChatModel chatModel;

    @GetMapping(value = "/generateStream", produces = "text/html;charset=utf-8")
    public Flux<String> generateStream(@RequestParam(value = "message", defaultValue = "你是谁") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));

        // 原子布尔值跟踪分隔线状态
        AtomicBoolean needSeparate = new AtomicBoolean(true);

        return chatModel.stream(prompt)
                .mapNotNull(chatResponse -> {
                    // 获取响应内容
                    DeepSeekAssistantMessage deepSeekAssistantMessage = (DeepSeekAssistantMessage) chatResponse.getResult().getOutput();
                    // 推理内容
                    String reasoningContent = deepSeekAssistantMessage.getReasoningContent();
                    // 推理结果之后的正式回答
                    String text = deepSeekAssistantMessage.getText();

                    // 是否正式回答
                    boolean isTextResponse = false;
                    String rawContent;
                    if (Objects.isNull(text)) {
                        rawContent = reasoningContent;
                    } else {
                        rawContent = text;
                        isTextResponse = true; // 标记为正式回答
                    }
                    // 处理换行
                    String processed = StringUtils.isNotBlank(rawContent) ? rawContent.replace("\n", "<br>").replace("\r", "<br>") : rawContent;

                    if (isTextResponse
                            && needSeparate.compareAndSet(true, false)) {
                        processed = "<hr>" + processed;
                    }

                    return processed;
                });
    }
}
